#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

def getInfo(current_str):
    from BeautifulSoup import BeautifulSoup

    soup = BeautifulSoup(current_str)
    article_type = ""
    journal = ""
    journal_type = ""
    try:
        article_type = soup.find("article")["article-type"]
    except:
        article_type = ""
    try:
        journal = soup.find("journal-id", attrs={"journal-id-type": "nlm-ta"}).string
    except:
        journal = ""
    try:
        journal_type = soup.find("subj-group", attrs={"subj-group-type": "heading"}).subject.string
    except:
        journal_type = ""
    if article_type == None:
        article_type = ""
    if journal == None:
        journal = ""
    if journal_type == None:
        journal_type = ""
    return article_type, journal, journal_type


if __name__ == "__main__":
    article_types = {}
    journal_types = {}
    current_str = ""
    count = 0
    for i in xrange(1, 368):
        infile = open("pmc_id_" + str(i) + ".xml", 'r')
        infile.readline()
        infile.readline()
        current_str = infile.readline()[16:]
        for line in infile:
            line = line.rstrip()
            if not line.startswith("</pmc-articleset>"):
                current_str += line
                if line.startswith("</article>"):
                    current_str = current_str.strip()
                    if current_str != "":
                        count += 1
                        article_type, journal, journal_type = getInfo(current_str)
                        if article_type != "":
                            if article_types.has_key(article_type):
                                article_types[article_type] = article_types[article_type] + 1
                            else:
                                article_types[article_type] = 1
                        if journal != "" and journal_type != "":
                            if not journal_types.has_key(journal):
                                journal_types[journal] = {}
                            if journal_types[journal].has_key(journal_type):
                                journal_types[journal][journal_type] = journal_types[journal][journal_type] + 1
                            else:
                                journal_types[journal][journal_type] = 1
                        print str(count) + ":" + article_type + "|" + journal + "|" + journal_type
                    current_str = ""
        infile.close()

        outfile1 = open("type/article_types_"+str(i)+".txt", "w")
        for item in article_types.keys():
            outfile1.write(item + "\t" + str(article_types[item]) + "\n")
            outfile1.flush()
        outfile1.close()

        outfile2 = open("type/journal_types_"+str(i)+".txt", "w")
        for item in journal_types.keys():
            val = ""
            for item2 in journal_types[item].keys():
                val += item2 + "=" + str(journal_types[item][item2]) + "|"
            val = val[:-1]
            outfile2.write(item + "####" + val + "\n")
            outfile2.flush()
        outfile2.close()

        article_types.clear()
        journal_types.clear()
